import dolphindb as ddb
import pandas as pd
from datetime import datetime, timedelta

class SpreadProvider:
    """期货价差数据提供者"""
    
    def __init__(self):
        self.session = None
        
    def connect_db(self):
        """连接DolphinDB数据库"""
        if self.session is None:
            self.session = ddb.session()
            self.session.connect("localhost", 8848, "admin", "123456")
            
    def close_db(self):
        """关闭数据库连接"""
        if self.session:
            self.session.close()
            self.session = None
            
    def get_spread_data(self, code1, code2, start_date=None, end_date=None):
        """获取两个合约的价差数据
        
        Parameters
        ----------
        code1 : str
            第一个合约代码
        code2 : str
            第二个合约代码
        start_date : str, optional
            起始日期，格式：'YYYY-MM-DD'
        end_date : str, optional
            结束日期，格式：'YYYY-MM-DD'
            
        Returns
        -------
        pd.DataFrame
            包含以下列的DataFrame：
            - date: 日期
            - code1: 第一个合约代码
            - code2: 第二个合约代码
            - spread: 价差（code1收盘价 - code2收盘价）
            - code1_close: 第一个合约收盘价
            - code2_close: 第二个合约收盘价
            - code1_volume: 第一个合约成交量
            - code2_volume: 第二个合约成交量
        """
        try:
            self.connect_db()
            
            # 构建查询条件
            conditions = f"code in [`{code1}, `{code2}]"
            if start_date:
                conditions += f" and date >= date('{start_date}')"
            if end_date:
                conditions += f" and date <= date('{end_date}')"
            
            # 查询数据
            script = f"""
            select date, code, close, volume
            from loadTable("dfs://futures", "daily_all")
            where {conditions}
            order by date
            """
            
            df = self.session.run(script)
            
            if df.empty:
                return None
                
            # 将数据分为两个合约
            df1 = df[df['code'] == code1].copy()
            df2 = df[df['code'] == code2].copy()
            
            # 确保两个合约在同一天都有数据
            merged_df = pd.merge(df1, df2, on='date', suffixes=('_1', '_2'))
            
            # 计算价差
            result = pd.DataFrame({
                'date': merged_df['date'],
                'code1': code1,
                'code2': code2,
                'spread': merged_df['close_1'] - merged_df['close_2'],
                'code1_close': merged_df['close_1'],
                'code2_close': merged_df['close_2'],
                'code1_volume': merged_df['volume_1'],
                'code2_volume': merged_df['volume_2']
            })
            
            return result
            
        except Exception as e:
            print(f"获取价差数据失败: {str(e)}")
            return None
        finally:
            self.close_db()
            
    def get_continuous_spread(self, product1, product2, months1, months2, start_date=None, end_date=None):
        """获取连续价差数据
        
        Parameters
        ----------
        product1 : str
            第一个品种代码（如 'RB'）
        product2 : str
            第二个品种代码（如 'HC'）
        months1 : list
            第一个品种的合约月份列表（如 [1, 5, 10]）
        months2 : list
            第二个品种的合约月份列表（如 [1, 5, 10]）
        start_date : str, optional
            起始日期，格式：'YYYY.MM.DD'
        end_date : str, optional
            结束日期，格式：'YYYY.MM.DD'
            
        Returns
        -------
        pd.DataFrame
            连续价差数据
        """
        # TODO: 实现连续价差计算逻辑
        pass

def main():
    """测试代码"""
    provider = SpreadProvider()
    
    # 测试获取RB2401和HC2401的价差
    df = provider.get_spread_data('l2401', 'v2401', '2023.01.01', '2024.01.01')
    
    if df is not None:
        print("\n价差数据示例:")
        print(df.head())
        
        # 保存到csv文件（可选）
        # df.to_csv('spread_data.csv', index=False)

if __name__ == "__main__":
    main() 